import argparse
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from networks.vision_transformer import SwinUnet as ViT_seg
from networks.swin_transformer_unetPP import SwinTansformerUnetPP as SwinUnetPP
from networks.unetPP import UNetPP as UnetPP
from networks.unet import UNet as UNet
from networks.u2net import u2net_lite as U2net
from networks.vgg_unet import VGG16UNet as VGG16UNet
from networks.mobilenet_unet import MobileV3Unet as MobileUnet
from trainer import trainer_synapse
from config import get_config

parser = argparse.ArgumentParser()
# 1.基本参数设置
# 指定数据集的根目录
# 使用 type=str 指定参数值的类型为字符串，默认值为 '../Synapse/train_npz'，帮助信息为 'Name of Experiment'
parser.add_argument('--root_path', type=str,
                    default='../data/Synapse/train_npz', help='root dir for data')
# 这段代码的作用是添加参数选项，用于指定实验名称
# 使用 type=str 指定参数值的类型为字符串，默认值为 'Synapse'，帮助信息为 'experiment_name'
parser.add_argument('--dataset', type=str,
                    default='Synapse', help='experiment_name')
# 这段代码的作用是确定list文件的位置
parser.add_argument('--list_dir', type=str,
                    default='./lists/Synapse', help='list dir')
# 这段代码的作用是添加参数选项，用于指定输出通道数
# 使用 type=int 指定参数值的类型为整数，默认值为 9，帮助信息为 'output channel of network'
# 输出通道的意思是网络输出的特征图的通道数，定义网络的输出通道数，对应分割任务的类别数量
parser.add_argument('--num_classes', type=int,
                    default=9, help='output channel of network')
# 输出文件的位置
parser.add_argument('--output_dir', type=str, help='output dir')
# 这段代码的作用是添加参数选项，用于指定最大迭代次数
# 使用 type=int 指定参数值的类型为整数，默认值为 30000，帮助信息为 'maximum epoch number to train'
parser.add_argument('--max_iterations', type=int,
                    default=30000, help='maximum epoch number to train')
# 设定最大的epochs
parser.add_argument('--max_epochs', type=int,
                    default=150, help='maximum epoch number to train')
# 这段代码的作用是添加参数选项，用于指定批量大小
# 使用 type=int 指定参数值的类型为整数，默认值为 24，帮助信息为 'batch_size per gpu'
parser.add_argument('--batch_size', type=int,
                    default=24, help='batch_size per gpu')
parser.add_argument('--n_gpu', type=int, default=1, help='total gpu')
# 这段代码的作用是添加参数选项，用于指定是否使用确定性训练（用于可重复性）
# 使用 type=int 指定参数值的类型为整数，默认值为 1，帮助信息为 'whether use deterministic training'
parser.add_argument('--deterministic', type=int, default=1,
                    help='whether use deterministic training')
# 这段代码的作用是添加参数选项，用于指定学习率
# 使用 type=float 指定参数值的类型为浮点数，默认值为 0.01，帮助信息为 'segmentation network learning rate'
parser.add_argument('--base_lr', type=float, default=0.01,
                    help='segmentation network learning rate')
# 图像大小设置
parser.add_argument('--img_size', type=int,
                    default=224, help='input patch size of network input')
# 这段代码的作用是添加参数选项，用于指定随机种子，确保实验可重复性
# 使用 type=int 指定参数值的类型为整数，默认值为 3407，帮助信息为 'random seed'
parser.add_argument('--seed', type=int,
                    default=3407, help='random seed')
# 配置文件路径
parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', )
# 配置文件额外参数
parser.add_argument(
    "--opts",
    help="Modify config options by adding 'KEY VALUE' pairs. ",
    default=None,
    nargs='+',
)


# 用zip文件而不是文件夹
parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
# 数据缓存方式
parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
                    help='no: no cache, '
                         'full: cache all data, '
                         'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
# 从检查点恢复训练
parser.add_argument('--resume', help='resume from checkpoint')
# 梯度加速步骤
parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
# 是否启动梯度检查
parser.add_argument('--use-checkpoint', action='store_true',
                    help="whether to use gradient checkpointing to save memory")
parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'],
                    help='mixed precision opt level, if O0, no amp is used')
# 别名
parser.add_argument('--tag', help='tag of experiment')
# 只评估不训练
parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
# 只测试不训练
parser.add_argument('--throughput', action='store_true', help='Test throughput only')



# parser.add_argument("--dataset_name", default="datasets")
# 类别数
parser.add_argument("--n_class", default=9, type=int)
# 数据载入线程
parser.add_argument("--num_workers", default=8, type=int)
# 内部评估参数
parser.add_argument("--eval_interval", default=1, type=int)

# 将参数导入配置项,更新后返回全部模型参数
args = parser.parse_args()
if args.dataset == "Synapse":
    args.root_path = os.path.join(args.root_path, "train_npz")
config = get_config(args)

if __name__ == "__main__":
    # CUDNN 配置
    if not args.deterministic:
        cudnn.benchmark = True
        cudnn.deterministic = False
    else:
        cudnn.benchmark = False
        cudnn.deterministic = True

    # 设置随机种子,确保实验结果的可复现性。
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    # 配置数据集路径和类别
    dataset_name = args.dataset
    dataset_config = {
        args.dataset: {
            'root_path': args.root_path,
            'list_dir': f'./lists/{args.dataset}',
            'num_classes': args.n_class,
        },
    }
    # 如果 batch_size 不是 24，但它是 6 的倍数，代码会根据当前的批量大小调整学习率(正比关系)。
    if args.batch_size != 24 and args.batch_size % 6 == 0:
        args.base_lr *= args.batch_size / 24
    # 数据集的类别数。
    args.num_classes = dataset_config[dataset_name]['num_classes']
    # 数据集的根路径。
    args.root_path = dataset_config[dataset_name]['root_path']
    # 数据集列表的路径。
    args.list_dir = dataset_config[dataset_name]['list_dir']


    # 创建输出目录
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # vit
    # 模型初始化
    net = ViT_seg(config, img_size=args.img_size, num_classes=args.num_classes).cuda()
    # 预训练初始化
    net.load_from(config)
    # unet++
    # net = UnetPP().cuda()
    # unet
    # net = UNet(n_channels=1 , n_classes=9).cuda()
    # u2net
    # net = U2net().cuda()
    # vgg_unet

    # mobilenet
    # net = MobileUnet(num_classes=9).cuda()
    # 开始训练
    # trainer = {'Synapse': trainer_synapse}
    trainer_synapse(args, net, args.output_dir)


# python train.py --output_dir ./model_out/datasets --dataset datasets --img_size 224 --batch_size 32 --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path /media/aicvi/11111bdb-a0c7-4342-9791-36af7eb70fc0/NNUNET_OUTPUT/nnunet_preprocessed/Dataset001_mm/nnUNetPlans_2d_split